#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <bitset>
#include <cmath>
#include "dct.h"
#include "encode.h"
#include "common.h"

float * readImg(int &dim1,int &dim2,std::string filename)
{
    std::ifstream fs(filename.c_str(),std::fstream::in|std::fstream::binary);
    std::string type;
    std::string dim;
    std::string col;
    std::getline(fs,type);
    std::getline(fs,dim);
    std::getline(fs,col);
    std::istringstream ss(dim);
    ss>>dim1>>dim2;
    float *img=new float[dim1*dim2];
    for(int i=0;i!=dim1;++i)
    {
        for (int j=0;j!=dim2;++j)
        {
            *(img+i*dim1+j)=(float)fs.get();
        }
    }
    fs.close();
    return img;
}
int **quantize(float **blk, int n[BLK_DIM][BLK_DIM], float qs)
{
    int **quantized_mat=new int* [BLK_DIM];
    for (int i=0;i!=BLK_DIM;++i) quantized_mat[i]=new int [BLK_DIM];
    for (int i=0;i!=BLK_DIM;++i)
        for (int j=0;j!=BLK_DIM;++j)
            quantized_mat[i][j]=(int)(blk[i][j]/n[i][j]/qs+0.5);
    return quantized_mat;
}
int *scan(int **quantized_mat)
{
    int scanM[BLK_DIM][BLK_DIM]=
        {   {0,  1,  5,  6,  14, 15, 27, 28},
            {2,  4,  7,  13, 16, 26, 29, 42},
            {3,  8,  12, 17, 25, 30, 41, 43},
            {9,  11, 18, 24, 31, 40, 44, 53},
            {10, 19, 23, 32, 39, 45, 52, 54},
            {20, 22, 33, 38, 46, 51, 55, 60},
            {21, 34, 37, 47, 50, 56, 59, 61},
            {35, 36, 48, 49, 57, 58, 62, 63}};
    int *vec1D=new int[BLK_DIM*BLK_DIM];
    for (int i=0;i!=BLK_DIM;++i)
        for (int j=0;j!=BLK_DIM;++j)
            vec1D[scanM[i][j]]=quantized_mat[i][j];
    return vec1D;
}
void writeBitSeq(std::vector<bool> &bitbuffer, std::vector<bool> &bitSeq)
{
    for (std::vector<bool>::iterator it=bitSeq.begin();it!=bitSeq.end();++it)
    {
        //std::cout<<(*it);
        bitbuffer.push_back(*it);
    }
}
int getCat(int value)
{
    if(value==0)return 0;
    if(value>0) return floor(log2(value)+1);
    if(value<0) return floor(log2(-value)+1);
}
void writeAmp(int size, int value,std::vector<bool> &bitbuffer)
{
    bool isPositive=(value>0);
    for (int i=0;i!=size;++i)
    {
        bool binaryDigit=value%(int)pow(2,size-i)/(int)pow(2,size-i-1);
        binaryDigit=isPositive?binaryDigit:!binaryDigit;    //if the number is negative, do one's complement
        //std::cout<<binaryDigit;
        bitbuffer.push_back(binaryDigit);
    }
}



void huffmanDC(std::vector<bool> &bitbuffer,int delta_dc,std::map<int,std::vector<bool> > &codebook)
{
    int cat;
    bool isPositive;
    cat=getCat(delta_dc);    
    //std::cout<<"CAT: "<<cat<<std::endl;
    writeBitSeq(bitbuffer,codebook[cat]);
    writeAmp(cat,delta_dc,bitbuffer);
}
void huffmanAC(std::vector<bool> &bitbuffer,int *vec1D,std::map<std::pair<int,int>,std::vector<bool> > &codebook)
{
    int countZero=0;
    std::pair<int,int> key;
    for (int i=1;i!=BLK_DIM*BLK_DIM;++i)
    {
        if (vec1D[i]==0)
        {
            countZero++;
            continue;
        }
        while (countZero>=15)
        {
            key=std::pair<int,int>(15,0);
            //std::cout<<"\tRUN: "<<15<<"\tCAT: "<<0<<"\t";
            countZero-=15;
            writeBitSeq(bitbuffer,codebook[key]);        
            //std::cout<<std::endl;
        }
        int cat=getCat(vec1D[i]);
        //std::cout<<"\tRUN: "<<countZero<<"\tCAT: "<<cat<<"\t";
        key=std::pair<int,int>(countZero,cat);
        writeBitSeq(bitbuffer,codebook[key]);
        writeAmp(cat,vec1D[i],bitbuffer);
        //std::cout<<std::endl;
        countZero=0;
    }
    if(countZero!=0)
    {
        key=std::pair<int,int>(0,0);
        //std::cout<<"\tRUN: "<<0<<"\tCAT: "<<0<<"\t";
        writeBitSeq(bitbuffer,codebook[key]);
        //std::cout<<std::endl;
    }
}
int charToInt(char a)
{
    switch(a)
    {
        case 'A':return 10;
        case 'B':return 11;
        case 'C':return 12;
        case 'D':return 13;
        case 'E':return 14;
        case 'F':return 15;
        default:
                 int x;
                 std::stringstream s;
                 s<<a;
                 s>>x;
                 return x;
    }
}

void load_codebook(std::map<int,std::vector<bool> > &dc_codebook,std::map<std::pair<int,int>,std::vector<bool> > &ac_codebook)
{
    std::ifstream codebookfile(CODEBOOK);
    std::string line;
    int key;
    int subKey;
    int count=0;
    std::vector<bool> value;
    while(std::getline(codebookfile,line))
    {
        value.clear();
        if (line=="")continue;
        key=charToInt(line[0]);
        if(count<12){
            for(int i=1;i!=line.length();++i)
            {
                if(line[i]=='0') value.push_back(0);
                if(line[i]=='1') value.push_back(1);
            }
            dc_codebook.insert(std::pair<int, std::vector<bool> >(key,value));
        }
        else
        {
            subKey=charToInt(line[2]);
            std::pair<int,int> ac_key=std::pair<int,int>(key,subKey);
            for (int i=3;i!=line.length();++i)
            {
                if(line[i]=='0') value.push_back(0);
                if(line[i]=='1') value.push_back(1);
            }
            ac_codebook.insert(std::pair<std::pair<int,int>, std::vector<bool> >(ac_key,value));
        }
        count++;
    }
    codebookfile.close();
}
void write_jpg(std::string filename,std::vector<bool> &bitbuffer)
{
    std::ofstream fs(filename.c_str(),std::fstream::out|std::fstream::binary);
    char byte;
    int i=0;
    for (unsigned int i=0;i<bitbuffer.size();++i)
    {
        if(i%8==0&&i!=0)
        {
            fs.write(&byte,1);
        }
        if(bitbuffer[i])
        {
            byte |= (1<<(7-(i%8)));
        }
        else
        {
            byte &= ~(1<<(7-(i%8)));
        }
    }
    fs.write(&byte,1);
    fs.close();
}
int encode(float qs,int &nbits)
{
    int dim1,dim2;                                  //dimension of original img file
    float *img;
    img=readImg(dim1,dim2,ORIGINAL_IMG);

    for (int i=0;i!=dim1*dim2;++i)
    {
        *(img+i)-=128;
    }
    dct(img,dim1,8,1);
 
    std::map<int,std::vector<bool> > dc_codebook;
    std::map<std::pair<int,int>,std::vector<bool> > ac_codebook;
    load_codebook(dc_codebook,ac_codebook);
    
    /*
    std::cout<<"DC codebook"<<std::endl;
    for (std::map<int,std::vector<bool> >::iterator it=dc_codebook.begin();it!=dc_codebook.end();++it)
    {
        std::cout<<it->first<<" : ";
        for (std::vector<bool>::iterator it2=(it->second).begin();it2!=(it->second).end();++it2)
        {
            std::cout<<*(it2);
        }
        std::cout<<std::endl;
    }

    std::cout<<"AC codebook"<<std::endl;
    for (std::map<std::pair<int,int>,std::vector<bool> >::iterator it=ac_codebook.begin();it!=ac_codebook.end();++it)
    {
        std::cout<<std::get<0>(it->first)<<" "<<std::get<1>(it->first)<<" : ";
        for (std::vector<bool>::iterator it2=(it->second).begin();it2!=(it->second).end();++it2)
        {
            std::cout<<*(it2);
        }
        std::cout<<std::endl;
    }
    */

    float **blk;
    int blky, blkx, pby, pbx;
    int **quantized_mat;
    int *vec1D;
    int prevDC;
    int dc;
    int delta_dc;
    bool isFirstDC=true;
    std::vector<bool> bitbuffer;
    int n[BLK_DIM][BLK_DIM]={{16,11,10,16,24,40,51,61},
                                    {12,12,14,19,26,58,60,55},
                                    {14,13,16,24,40,57,69,59},
                                    {14,17,22,29,51,87,80,62},
                                    {18,22,37,56,68,109,103,77},
                                    {24,35,55,64,81,104,113,92},
                                    {49,64,78,87,103,121,120,101},
                                    {72,92,95,98,112,110,103,99}};

    blk=new float *[BLK_DIM];

    for (int i=0;i!=BLK_DIM;++i)
    {
        blk[i]=new float[BLK_DIM];
    }

    for(blky=0; blky<dim2/BLK_DIM; blky++) {
        for(blkx=0; blkx<dim1/BLK_DIM; blkx++) {
            // get one block data
            for(pby=0;pby<BLK_DIM;pby++)
                for(pbx=0;pbx<BLK_DIM;pbx++)
                    blk[pby][pbx]=img[(blky*BLK_DIM+pby)*dim1+(blkx*BLK_DIM+pbx)];
            quantized_mat=quantize(blk,n,qs);
            /*
            std::cout<<"QUANTIZED MAT"<<std::endl;
            for (int i=0;i!=BLK_DIM;++i)
            {
                for (int j=0;j!=BLK_DIM;++j)
                {
                    std::cout<<quantized_mat[i][j]<<" ";
                }
                std::cout<<std::endl;
            }
            std::cout<<"====================================="<<std::endl;
            */
            vec1D=scan(quantized_mat);
            dc=vec1D[0];
            delta_dc=isFirstDC?dc:(dc-prevDC);
            prevDC=dc;
            isFirstDC=false;
            //std::cout<<std::endl<<"=========================="<<std::endl;
            //std::cout<<"DC: ";
            huffmanDC(bitbuffer,delta_dc,dc_codebook);
            //std::cout<<std::endl<<"---------------------------"<<std::endl;
            //std::cout<<"AC: ";
            huffmanAC(bitbuffer,vec1D,ac_codebook);
            //std::cout<<std::endl<<"=========================="<<std::endl;
        }
    }
    nbits=bitbuffer.size();
    write_jpg(COMPRESSED_IMG,bitbuffer);

    for (int i=0;i!=BLK_DIM;++i)
    {
        delete[] blk[i];
        delete[] quantized_mat[i];
    }
    delete[] blk;
    delete[] quantized_mat;
    delete[] img;
    delete[] vec1D;

    /* testing */

    return 0;
}

